import torch
import torch.nn as nn
import torch.nn.functional as F

class EncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, nhead, dropout=dropout)
        # 实现一个线性变换和一个dropout层，用于多头注意力层的输出
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.feedforward = PositionwiseFeedForward(d_model, dim_feedforward, dropout)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        r"""Pass the input through the encoder layer.

        Args:
            src: Tensor, shape [seq_len, batch, embed_dim]
            src_mask: Tensor, shape [seq_len, seq_len], optional,
                The mask for the src sequence. Optional masked encoding layers
                with masked positions.
            src_key_padding_mask: Tensor, shape [batch, seq_len], optional,
                A byte tensor where positions are set to `True` for padded
                positions in `src`.

        Returns:
            output: Tensor, shape [seq_len, batch, embed_dim]
        """
        src2 = self.norm1(src + self.dropout1(self.self_attn(src, src, src, attn_mask=src_mask,
                                                             key_padding_mask=src_key_padding_mask)[0]))
        src = self.norm2(src2 + self.dropout2(self.feedforward(src2)))
        return src

class Encoder(nn.Module):
    def __init__(self, encoder_layers, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(Encoder, self).__init__()
        self.layers = nn.ModuleList([EncoderLayer(d_model, nhead, dim_feedforward, dropout)
                                    for _ in range(encoder_layers)])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, src, mask=None, src_key_padding_mask=None):
        r"""Pass the input through the encoder layers in turn.

        Args:
            src: Tensor, shape [seq_len, batch, embed_dim]
            mask: optional tensor, shape [seq_len, seq_len]
            src_key_padding_mask: Tensor, shape [batch, seq_len], optional,
                A byte tensor where positions are set to `True` for padded
                positions in `src`.

        Returns:
            output: Tensor, shape [seq_len, batch, embed_dim]
        """
        output = src
        
        for layer in self.layers:
            output = layer(output, src_mask=mask,
                           src_key_padding_mask=src_key_padding_mask)

        return self.norm(output)

